from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.prebuilt import chat_agent_executor

from model import get_model

# sqlalchemy, 初始化 MySql 数据库连接
HOST = '192.168.110.18'
PORT = '3306'
DATABASE = 'yg-llm-demo'
USER = 'root'
PASSWORD = 'root'

MYSQL_URI = 'mysql+pymysql://{user}:{password}@{host}:{port}/{db}?charset=utf8'.format(
    user=USER, password=PASSWORD, host=HOST, port=PORT, db=DATABASE)

# 1. 连接 MySQL 数据库
db = SQLDatabase.from_uri(MYSQL_URI)

# 2. 调用大模型
zhipuai = get_model()

# 创建工具
toolkit = SQLDatabaseToolkit(db=db, llm=zhipuai)
tools = toolkit.get_tools()

# 提示词
system_prompt = """
你是一个被设计用来与SQL数据库交互的代理。
将给定的输入问题，创建一个语法正确的SQL语句并执行，然后查看查询结果并返回答案。

你可以使用与数据库交互的工具。在执行查询之前，你必须仔细检查。如果在执行查询时出现错误，请重写查询并重试。
不要对数据库做任何DML语句（插入，更新，删除等）

首先，你应该查看数据库中的表，看看可以查询什么。不要跳过这一步。
然后查询最相关的表的模式。
"""
system_message = SystemMessage(content=system_prompt)

# 创建代理
agent_executor = chat_agent_executor.create_tool_calling_executor(zhipuai, tools,messages_modifier=system_message)

resp = agent_executor.invoke({'messages': [HumanMessage(content='请问Alice的电子邮件是多少？')]})

result = resp['messages']
print(f"result:{result}")
# print(len(result))

print(result[-1].content)  # 打印最后一条消息的内容
